import os
import sys
import json
from tqdm import tqdm



from .make_labels_uk import MakeLabelsUK
from ..utils.processing_uk import UKBioBankProcessor


class CheckLabelFilesUK:
    def __init__(self, check_path: str, expected_line_count: int, regenerate: bool = False, **kwargs):
        self.check_path = check_path
        self.expected_line_count = expected_line_count
        self.regenerate = regenerate
        self.build_labels_instance = None
        if self.regenerate:
            save_path = os.path.dirname(check_path)
            self.build_labels_instance = MakeLabelsUK(
                dataset_path=kwargs['dataset_path'],
                save_path=save_path,
                id_file_path=kwargs['id_file_path'],
                encoding_file_path=kwargs['encoding_file_path'],
                instance_file_path=kwargs['instance_file_path'],
                categories_file_path=kwargs['categories_file_path'],
                api_key=kwargs.get('api_key', "")
            )
            data_processor = UKBioBankProcessor(kwargs['id_file_path'], kwargs['encoding_file_path'],
                                            kwargs['instance_file_path'], kwargs['categories_file_path'])
            self.build_labels_instance.fields = data_processor.fields
            self.build_labels_instance.instances = data_processor.instances
            self.build_labels_instance.categories = data_processor.categories
            self.build_labels_instance.save_labels_list_dir = os.path.join(save_path, "labels")


    def check_files(self):
        if not os.path.exists(self.check_path):
            print(f"Error: The path '{self.check_path}' does not exist.")
            return

        print(f"Checking files in: {self.check_path}")
        print(f"Expected line count: {self.expected_line_count}")

        mismatched_files = []

        for filename in tqdm(os.listdir(self.check_path), desc="Checking files"):
            if filename.endswith(".json"):
                file_path = os.path.join(self.check_path, filename)
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                        actual_line_count = len(data)
                        if actual_line_count != self.expected_line_count:
                            mismatched_files.append((filename, actual_line_count))
                except json.JSONDecodeError:
                    print(f"Error decoding JSON from file: {filename}")
                    mismatched_files.append((filename, -1))
                except Exception as e:
                    print(f"Error processing file {filename}: {e}")

        if mismatched_files:
            print("\n--- Mismatched Files ---")
            for filename, actual_line_count in mismatched_files:
                if actual_line_count == -1:
                    print(f"File: {filename}, Error reading or JSON decoding error.")
                else:
                    print(f"File: {filename}, Actual Lines: {actual_line_count}, Expected: {self.expected_line_count}")
            print(f"\nTotal mismatched files: {len(mismatched_files)}")

            if self.regenerate:
                print("\n--- Regenerating Mismatched Files ---")
                regenerate_list = [f for f, l in mismatched_files]
                for filename in tqdm(regenerate_list, desc="Regenerating files"):
                    column_name = filename[:-5]  # remove .json
                    self.build_labels_instance.process_column_name(column_name)
                print("\nRegeneration complete.")
        else:
            print("\nAll files have the expected number of lines.")

    def main(self):
        self.check_files()


if __name__ == '__main__':
    # Configuration for checking and regenerating files
    check_path = os.path.join(project_dir_path, "datasets/ready/ukbiobank/v3/labels")
    expected_lines = 336
    regenerate = True  # Set to True to enable regeneration of mismatched files

    # Parameters for BuildLabels, required if regenerate is True
    dataset_path = "/home/xw/python3test/omics-bert-datasets/datasets/ready/ukbiobank/v4"
    id_file_path = os.path.join(project_dir_path, "datasets/raw/ukbiobank/v3/id.csv")
    encoding_file_path = os.path.join(project_dir_path, "datasets/raw/ukbiobank/v3/encoding.csv")
    instance_file_path = os.path.join(project_dir_path, "datasets/raw/ukbiobank/v3/instance.csv")
    categories_file_path = os.path.join(project_dir_path, "datasets/raw/ukbiobank/v3/category.csv")
    api_key = ""

    checker = CheckLabelFilesUK(
        check_path=check_path,
        expected_line_count=expected_lines,
        regenerate=regenerate,
        dataset_path=dataset_path,
        id_file_path=id_file_path,
        encoding_file_path=encoding_file_path,
        instance_file_path=instance_file_path,
        categories_file_path=categories_file_path,
        api_key=api_key
    )
    checker.main()